// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/expr.h>
#include <pollux/expression/lambda_expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/lambda_function_util.h>
#include <pollux/functions/lib/rows_translation_util.h>
#include <pollux/vector/function_vector.h>

namespace kumo::pollux::functions {
    namespace {
        class FilterFunctionBase : public exec::VectorFunction {
        protected:
            // Applies filter functions to elements of maps or arrays and returns the
            // number of elements that passed the filters. Stores the number of elements
            // in each array or map that passed the filter in resultSizes. Stores the
            // indices of passing elements in selectedIndices. resultOffsets is populated
            // assuming sequential layout of the passing elements.
            template<typename T>
            static vector_size_t doApply(
                const SelectivityVector &rows,
                const std::shared_ptr<T> &input,
                const VectorPtr &lambdas,
                const std::vector<VectorPtr> &lambdaArgs,
                exec::EvalCtx &context,
                BufferPtr &resultOffsets,
                BufferPtr &resultSizes,
                BufferPtr &selectedIndices) {
                const auto *inputOffsets = input->rawOffsets();
                const auto *inputSizes = input->rawSizes();

                auto *pool = context.pool();
                resultSizes = allocateSizes(rows.end(), pool);
                resultOffsets = allocateOffsets(rows.end(), pool);
                auto *rawResultSizes = resultSizes->asMutable<vector_size_t>();
                auto *rawResultOffsets = resultOffsets->asMutable<vector_size_t>();

                const auto numElements = lambdaArgs[0]->size();

                selectedIndices = allocate_indices(numElements, pool);
                auto *rawSelectedIndices = selectedIndices->asMutable<vector_size_t>();

                vector_size_t numSelected = 0;

                auto elementToTopLevelRows =
                        getElementToTopLevelRows(numElements, rows, input.get(), pool);

                exec::LocalDecodedVector bitsDecoder(context);
                auto iter = lambdas->as_unchecked<FunctionVector>()->iterator(&rows);
                while (auto entry = iter.next()) {
                    auto elementRows =
                            toElementRows<T>(numElements, *entry.rows, input.get());
                    auto wrapCapture =
                            toWrapCapture<T>(numElements, entry.callable, *entry.rows, input);

                    VectorPtr bits;
                    entry.callable->apply(
                        elementRows,
                        nullptr,
                        wrapCapture,
                        &context,
                        lambdaArgs,
                        elementToTopLevelRows,
                        &bits);
                    bitsDecoder.get()->decode(*bits, elementRows);
                    entry.rows->applyToSelected([&](vector_size_t row) {
                        if (input->is_null_at(row)) {
                            return;
                        }
                        auto size = inputSizes[row];
                        auto offset = inputOffsets[row];
                        rawResultOffsets[row] = numSelected;
                        for (auto i = 0; i < size; ++i) {
                            if (!bitsDecoder.get()->is_null_at(offset + i) &&
                                bitsDecoder.get()->value_at<bool>(offset + i)) {
                                ++rawResultSizes[row];
                                rawSelectedIndices[numSelected] = offset + i;
                                ++numSelected;
                            }
                        }
                    });
                }

                selectedIndices->setSize(numSelected * sizeof(vector_size_t));

                return numSelected;
            }
        };

        // See documentation at
        //    - https://prestodb.io/docs/current/functions/array.html
        //    - https://prestodb.io/docs/current/functions/lambda.html
        //    - https://prestodb.io/blog/2020/03/02/presto-lambda
        class ArrayFilterFunction : public FilterFunctionBase {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2);
                exec::LocalDecodedVector arrayDecoder(context, *args[0], rows);
                auto &decodedArray = *arrayDecoder.get();

                auto flatArray = flattenArray(rows, args[0], decodedArray);

                VectorPtr elements = flatArray->elements();
                BufferPtr resultSizes;
                BufferPtr resultOffsets;
                BufferPtr selectedIndices;
                auto numSelected = doApply(
                    rows,
                    flatArray,
                    args[1],
                    {elements},
                    context,
                    resultOffsets,
                    resultSizes,
                    selectedIndices);

                // Filter can pass along very large elements vectors that can hold onto
                // memory and copy operations on them can further put memory pressure. We
                // try to flatten them if the dictionary layer is much smaller than the
                // elements vector.
                auto wrappedElements = numSelected
                                           ? BaseVector::wrap_in_dictionary(
                                               BufferPtr(nullptr),
                                               std::move(selectedIndices),
                                               numSelected,
                                               std::move(elements),
                                               true /*flattenIfRedundant*/)
                                           : nullptr;
                // Set nulls for rows not present in 'rows'.
                BufferPtr newNulls = addNullsForUnselectedRows(flatArray, rows);
                auto localResult = std::make_shared<ArrayVector>(
                    flatArray->pool(),
                    flatArray->type(),
                    std::move(newNulls),
                    rows.end(),
                    std::move(resultOffsets),
                    std::move(resultSizes),
                    wrappedElements);
                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // array(T), function(T, boolean) -> array(T)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("array(T)")
                    .argumentType("array(T)")
                    .argumentType("function(T,boolean)")
                    .build()
                };
            }
        };

        // See documentation at
        //    - https://prestodb.io/docs/current/functions/map.html
        //    - https://prestodb.io/docs/current/functions/lambda.html
        //    - https://prestodb.io/blog/2020/03/02/presto-lambda
        class MapFilterFunction : public FilterFunctionBase {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2);
                exec::LocalDecodedVector mapDecoder(context, *args[0], rows);
                auto &decodedMap = *mapDecoder.get();

                auto flatMap = flattenMap(rows, args[0], decodedMap);

                VectorPtr keys = flatMap->mapKeys();
                VectorPtr values = flatMap->mapValues();
                BufferPtr resultSizes;
                BufferPtr resultOffsets;
                BufferPtr selectedIndices;
                auto numSelected = doApply(
                    rows,
                    flatMap,
                    args[1],
                    {keys, values},
                    context,
                    resultOffsets,
                    resultSizes,
                    selectedIndices);

                // Filter can pass along very large elements vectors that can hold onto
                // memory and copy operations on them can further put memory pressure. We
                // try to flatten them if the dictionary layer is much smaller than the
                // elements vector.
                auto wrappedKeys = numSelected
                                       ? BaseVector::wrap_in_dictionary(
                                           BufferPtr(nullptr),
                                           selectedIndices,
                                           numSelected,
                                           std::move(keys),
                                           true /*flattenIfRedundant*/)
                                       : nullptr;
                auto wrappedValues = numSelected
                                         ? BaseVector::wrap_in_dictionary(
                                             BufferPtr(nullptr),
                                             selectedIndices,
                                             numSelected,
                                             std::move(values),
                                             true /*flattenIfRedundant*/)
                                         : nullptr;
                // Set nulls for rows not present in 'rows'.
                BufferPtr newNulls = addNullsForUnselectedRows(flatMap, rows);
                auto localResult = std::make_shared<MapVector>(
                    flatMap->pool(),
                    outputType,
                    std::move(newNulls),
                    rows.end(),
                    std::move(resultOffsets),
                    std::move(resultSizes),
                    wrappedKeys,
                    wrappedValues);
                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // map(K,V), function(K,V,boolean) -> map(K,V)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("map(K,V)")
                    .argumentType("map(K,V)")
                    .argumentType("function(K,V,boolean)")
                    .build()
                };
            }
        };
    } // namespace

    /// array_filter and map_filter are null preserving for the array and the
    /// map. But since an expr tree with a lambda depends on all named fields,
    /// including captures, a null in a capture does not automatically make a
    /// null result.

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_array_filter,
        ArrayFilterFunction::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<ArrayFilterFunction>());

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_map_filter,
        MapFilterFunction::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<MapFilterFunction>());
} // namespace kumo::pollux::functions
